from src import Preproduction as Pre

id, ch = Pre.load("results/hanzi_2_one_hot.data")
cnt = len(ch)
V = 10000

print("cnt = " + str(cnt))
print("V = " + str(V))

from src.EmbeddingGloVe import genX as genX
from src.EmbeddingGloVe import Glove as Glove

pathBase = 'dataset/chat_corpus/clean_chat_corpus/'
paths = [
    "chatterbot.tsv",
    "douban_single_turn.tsv",
    #"ptt.tsv",
    "qingyun.tsv",
    #"subtitle.tsv",
    "tieba.tsv",
    #"weibo.tsv"
]
fullPaths = [pathBase + path for path in paths]

import numpy as np
import random

from src.GeneratorLSTMv6 import LSTM

embDim = 128
hidDim = 512
seqLen = 30

lstm = LSTM(seqLen, embDim, hidDim, V)
lstm.load("results/V8LSTM050.model")

def Gen(cur):
    x = [id[i] if id[i] < V else 0 for i in cur]
    e = lstm.eval(x)
    res = ''
    for i in e:
        res += ch[i] if i > 1 else ''
    return res

while True:
    cur = input("Query> ")
    if len(cur) and cur[0] == '!' :
        code = cur.split()
        if (code[0] == '!exit'):
            break
        elif (code[0] == '!code'):
            print("  Res> " + str(id[code[1]]))
    else:
        print("  Gen> " + Gen(cur))
    